# -*- coding:utf8 -*-
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

matplotlib.use('TkAgg')
TOTAL_STEP = 10000


def get_idx_loss_lr_array_ms(path, step_end=TOTAL_STEP):
    res_np = []
    loss_list = []
    lr_list = []
    print("get_loss_and_index_array from ms:")
    with open(path, "r") as f:
        lines = f.readlines()
        for line in lines:
            if "overflow cond" in line:
                line_list = line.split(',')
                idx_r = line_list[3].index(":")
                loss_val = line_list[3][idx_r + 2: -1]
                loss_list.append(float(loss_val))
                lr_val = line_list[5].split(":")[-1]
                lr_list.append(float(lr_val))
                step_idx_r = line_list[2].index("/")
                step_idx_l = line_list[2].index("[")
                step = line_list[2][step_idx_l + 1: step_idx_r].lstrip(" ")
                if step == f"{step_end}":
                    break
                one_step1 = [int(step) - 1, loss_val, lr_val]
                one_step2 = [int(step), loss_val, lr_val]
                res_np.append(one_step1)
                res_np.append(one_step2)
    return np.array(res_np)


def main():
    filename = 'E:/电信AI_telechat/实验/端午-精度对齐/llama_bias_0608_lilei.txt'
    filename2 = 'E:/电信AI_telechat/实验/端午-精度对齐/telechat_0608_lilei.txt'
    ms_np = get_idx_loss_lr_array_ms(filename)
    ms_np2 = get_idx_loss_lr_array_ms(filename2)
    print(ms_np.shape)
    print(ms_np[:3, ])
    print(ms_np[-3:-1, ])
    loss_data = []
    loss_dict = {"llama_bias": ms_np, "telechat": ms_np2}
    labels = list(loss_dict.keys())
    min_data_len = min([len(loss) for loss in list(loss_dict.values())])
    print(min_data_len, labels)
    data_plot = []
    for loss in list(loss_dict.values()):
        print(loss[:, 1])
        data_plot.append(loss[:, 1])
    # data_plot = [loss[:, 1][:min_data_len] for loss in list(loss_data.values())]
    axis = np.arange(0, min_data_len)
    # for d in data_plot:
    #     plt.plot(axis, d)
    # plt.legend(labels=labels)
    # plt.xlabel("step")
    # plt.ylabel("loss")
    # plt.legend()
    # plt.grid(True)
    # plt.show()
    plt.figure(figsize=(10, 6))
    # l1 = np.arange(1, 10000, 0.2)
    # l2 = np.linspace(0.1, 0.3, 10000)
    # print(l1)
    # print(l2)
    # plt.plot(list(l1), label='llama_drop')
    # plt.plot(list(l2), label='llama_drop')
    a = [i for i in range(1000)]
    plt.plot(np.array(a), np.array(loss_dict["llama_bias"][:1000, 1]), label='llama_drop')
    # plt.plot(list(loss_dict["telechat"][:20, 1]), label='telechat')
    # plt.plot(loss1[:steps], loss2[:steps, 1], label='llama_bias')
    # plt.gca().invert_yaxis()
    plt.xlabel('step')
    plt.ylabel('loss')
    plt.title('loss info')
    plt.legend()
    plt.grid(True)
    plt.show()

    # plt.savefig(f'./loss.png', dpi=1000, bbox_inches='tight')


if __name__ == '__main__':
    main()
